# -*- coding: utf-8 -*-

#    OriginDB
#    Copyright (C) 2010  Jason Power <power.jg@gmail.com>
#
#    This program is free software; you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation; either version 2 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License along
#    with this program; if not, write to the Free Software Foundation, Inc.,
#    51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

import openbabel

from math import sqrt, pi


def genVariations(items, n):
	if n==0:
		yield []
	else:
		for i in xrange( len(items)-n+1):
			for v in genVariations(items[i+1:],n-1):
				yield [items[i]]+v


def degToRad(deg):
	return deg*pi/180


class MinAtom:
	def __init__(self, atom):
		self.atom = atom
		self.dx = 0
		self.dy = 0

		self.x = atom.GetX()
		self.y = atom.GetY()


class Minimize:

	def __init__(self):
		self.cutoff_rms_grad = 1e-5
		self.cutoff_max_grad = 3e-5
		self.max_iter_number = 1000

	def initialize(self, mol, bondLength=1, callback=None):
		self.mol = mol

		if bondLength < 0:
			bondLengths = []
			for bond in openbabel.OBMolBondIter(self.mol):
				l = bond.GetLength()  # COULD SEGFAULT!!!
				bondLengths.append(l)
			self.bondLength = sum(bondLengths)/len(bondLengths)
		else:
			self.bondLength = bondLength

		self.atoms = []
		for atom in openbabel.OBMolAtomIter(self.mol):
			self.atoms.append(MinAtom(atom))
		

	def optimizeCoords(self, bondLength=1, callback=None):

		sumDist, maxGrad = self.step()
		rmsGrad = sumDist / len(self.atoms)

		i=0
		ok = True
		self.endReason = "converged"

		while rmsGrad > (self.cutoff_rms_grad / self.bondLength) or maxGrad > (self.cutoff_max_grad / self.bondLength):
			i += 1
			if callback:
				callback(i, rmsGrad, maxGrad)

			sumDist, maxGrad = self.step()
			rmsGrad = sumDist / len(self.atoms)

			if rmsGrad > 1000:
				self.endReason = "gradien too large"
				ok = False
				break
			if i > self.max_iter_number:
				self.endReason = "max_iter_number iterations without converged result"
				ok = False
				break

		self.i = i
		self.rmsGrad = rmsGrad
		self.maxGrad = maxGrad
		return ok


	def finalize(self):
		for atom in self.atoms:
			atom.atom.SetVector(atom.x, atom.y, 0)


	def step(self, callback=None):

		for a in self.atoms:
			a.dx = 0
			a.dy = 0

		for a in self.atoms:
			self.setGradient(a)

		distSum = 0
		maxDist = 0
		for a in self.atoms:
			dx = a.dx
			dy = a.dy

			a.x += dx
			a.y += dy
			d = dx**2 + dy**2
			distSum += d

			if d > maxDist:
				maxDist = d

		return sqrt(distSum), sqrt(maxDist)


	def setGradient(self, atom):
		dx = 0
		dy = 0

		neighbors = self.getNeighbors(atom)
		for a in neighbors:
			gx, gy = self.getLengthGradient(atom, a)
			dx -= gx/2.0
			dy -= gy/2.0

		atom.dx += dx
		atom.dy += dy

		dx = 0
		dy = 0

		variations = list(genVariations(neighbors, 2))
		if variations:
			for var in variations:
				a1, a2 = var
				gx1, gy1, gx2, gy2 = self.getAngleGradient(degToRad(120), atom, a1, a2)
				a1.dx -= gx1/4.0
				a1.dy -= gy1/4.0
				a2.dx -= gx2/4.0
				a2.dy -= gy2/4.0


	def getLengthGradient(self, atom1, atom2, optLength=0):
		if not optLength:
			optLength = self.bondLength
		dx = atom1.x - atom2.x
		dy = atom1.y - atom2.y

		dist = sqrt(dx**2+dy**2)

		dd = dist - optLength

		gx = (dx / dist) * dd
		gy = (dy / dist) * dd

		return gx, gy


	def getAngleGradient(self, angle, refAtom, atom1, atom2):
		expl = sqrt(3) * self.bondLength
		gx1, gy1 = self.getLengthGradient(atom1, atom2, optLength=expl)
		gx2 = -gx1
		gy2 = -gy1

		return gx1, gy1, gx2, gy2


	def getNeighbors(self, atom):
		neighborAtomIdxs = []
		for bond in openbabel.OBMolBondIter(self.mol):
			if bond.GetBeginAtomIdx() == atom.atom.GetIdx():
				neighborAtomIdxs.append(bond.GetEndAtomIdx())
			elif bond.GetEndAtomIdx() == atom.atom.GetIdx():
				neighborAtomIdxs.append(bond.GetBeginAtomIdx())

		neighbors = []
		for id in neighborAtomIdxs:
			for atom in self.atoms:
				if atom.atom.GetIdx() == id:
					neighbors.append(atom)

		return neighbors